#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include "webrtc_vad.h"

static PyObject* vad_create(PyObject* self, PyObject* args) {
    VadInst* vad_inst = WebRtcVad_Create();
    if (vad_inst == NULL) {
        PyErr_SetString(PyExc_RuntimeError, "Failed to create VAD instance");
        return NULL;
    }
    return PyCapsule_New((void*)vad_inst, "VadInst", NULL);
}

static PyObject* vad_init(PyObject* self, PyObject* args) {
    PyObject* capsule;
    if (!PyArg_ParseTuple(args, "O", &capsule)) {
        return NULL;
    }
    VadInst* vad_instance = (VadInst*)PyCapsule_GetPointer(capsule, "VadInst");
    if (vad_instance == NULL) {
        return NULL;
    }
    if (WebRtcVad_Init(vad_instance) != 0) {
        PyErr_SetString(PyExc_RuntimeError, "Failed to initialize VAD instance");
        return NULL;
    }
    Py_RETURN_NONE;
}

static PyObject* vad_set_mode(PyObject* self, PyObject* args) {
    PyObject* capsule;
    int mode;
    if (!PyArg_ParseTuple(args, "Oi", &capsule, &mode)) {
        return NULL;
    }
    VadInst* vad_instance = (VadInst*)PyCapsule_GetPointer(capsule, "VadInst");
    if (vad_instance == NULL) {
        return NULL;
    }
    if (WebRtcVad_set_mode(vad_instance, mode) != 0) {
        PyErr_SetString(PyExc_RuntimeError, "Failed to set VAD mode");
        return NULL;
    }
    Py_RETURN_NONE;
}

static PyObject* vad_process(PyObject* self, PyObject* args) {
    PyObject* capsule;
    int sample_rate;
    const char* frame;
    Py_ssize_t frame_length;
    if (!PyArg_ParseTuple(args, "Ois#", &capsule, &sample_rate, &frame, &frame_length)) {
        return NULL;
    }
    VadInst* vad_instance = (VadInst*)PyCapsule_GetPointer(capsule, "VadInst");
    if (vad_instance == NULL) {
        return NULL;
    }
    int result = WebRtcVad_Process(vad_instance, sample_rate, (const int16_t*)frame, frame_length / 2);
    if (result < 0) {
        PyErr_SetString(PyExc_RuntimeError, "VAD process failed");
        return NULL;
    }
    return Py_BuildValue("i", result);
}

static PyMethodDef VadMethods[] = {
    {"create", vad_create, METH_VARARGS, "Create a new VAD instance"},
    {"init", vad_init, METH_VARARGS, "Initialize a VAD instance"},
    {"set_mode", vad_set_mode, METH_VARARGS, "Set the mode of a VAD instance"},
    {"process", vad_process, METH_VARARGS, "Process a frame with the VAD instance"},
    {NULL, NULL, 0, NULL}
};

static struct PyModuleDef vadmodule = {
    PyModuleDef_HEAD_INIT,
    "vad",
    NULL,
    -1,
    VadMethods
};

PyMODINIT_FUNC PyInit_vad(void) {
    return PyModule_Create(&vadmodule);
}